# -*- coding: utf-8 -*-
# @File:     rag_agent.py
# @Author:
# @DateTime: 2025/10/24/10:58
import datetime
import json
import time
import uuid
from typing import Optional, Union, Any, List, Dict
from fastapi.encoders import jsonable_encoder
from sqlalchemy import or_
from sqlalchemy.orm import Session

from app.api.chat.crud.chat import *
from app.api.chat.scheme.chat import *
from app.api.chat.utils.chat_util import ChatManage, RAGAgentManage
from app.db.database import get_db
from app.models.system import Users
from app.utils import security
from app.utils.custom_exc import CustomException
from app.utils.custom_response import success_response, error_response
from fastapi import APIRouter, Query, Body, Path, Depends, UploadFile, Form, File, Request


from langchain_core.messages import HumanMessage, ToolMessage, AIMessage
from sse_starlette.sse import EventSourceResponse

router = APIRouter()

chat_manage = RAGAgentManage()


@router.post("/rag_chat", summary="rag聊天生成", response_model=ChatResData, name='rag聊天生成')
async def chat(
    request: Request,
    *,
    db: Session = Depends(get_db),
    req_data: RAGChatReqData,
    current_user: Users = Depends(security.verify_access_token),
) -> Any:
    user_id = current_user.id
    dialog_uid = req_data.dialog_uid
    user_input = req_data.user_input
    file_id_list = req_data.file_list

    if dialog_uid:
        pass
    else:
        dialog_uid = uuid.uuid4().hex

    config = {"recursion_limit": 20, "configurable": {"thread_id": str(dialog_uid)}}

    history = list()
    if not chat_manage.checkpoint.get(config):
        create_history_message(dialog_uid, db, history)
    history.append(HumanMessage(content=user_input))

    if file_id_list:
        collection_list = select_file_collection_name(file_id_list, db, user_id)
    else:
        collection_list = []

    inputs = {"messages": history, "user_id": user_id, 'question': user_input, 'collection_list': collection_list, 'search_count': 0}
    await chat_manage.graph_build()

    data = {
        'dialog_uid': dialog_uid,
        'user_id': user_id,
        'agent_type': 'RAGAgent',
    }
    dialog_ins = create_dialog(data, db)

    dialog_ins_id = dialog_ins.id

    history_data = {
        'dialog_id': dialog_ins.id,
        'role': 'user',
        'content': user_input
    }
    create_dialog_history(history_data, db)

    def general_generate():
        data_res = ''
        try:
            data = {
                'data': '',
                'event': 'start',
                'dialog_uid': dialog_uid,
            }
            yield json.dumps(data, ensure_ascii=False)
            for event, response_metadata in chat_manage.graph.stream(inputs, config=config, stream_mode='messages'):

                # print(event)
                # print("-===", response_metadata)
                # print(event.response_metadata.get('finish_reason'))
                # 准备去调用工具
                if event.additional_kwargs.get('tool_calls'):

                    data = {
                        'data': '',
                        'event': 'tools',
                        'dialog_uid': dialog_uid,
                    }
                    yield json.dumps(data, ensure_ascii=False)
                    continue
                if event.response_metadata.get('finish_reason') == 'tool_calls':

                    data = {
                        'data': '',
                        'event': 'tools',
                        'dialog_uid': dialog_uid,
                    }
                    yield json.dumps(data, ensure_ascii=False)
                    continue
                # 调用
                if response_metadata['langgraph_node'] == 'tools':
                    data = {
                        'data': '',
                        'event': 'tools',
                        'dialog_uid': dialog_uid,
                    }
                    yield json.dumps(data, ensure_ascii=False)
                    continue
                if isinstance(event, ToolMessage):
                    data = {
                        'data': '',
                        'event': 'tools',
                        'dialog_uid': dialog_uid,
                    }
                    yield json.dumps(data, ensure_ascii=False)
                    continue
                data = {
                    'data': event.content,
                    'event': 'continue',
                    'dialog_uid': dialog_uid,
                }
                data_res += event.content
                yield json.dumps(data, ensure_ascii=False)
            data = {'data': {}, 'event': 'end', 'dialog_uid': dialog_uid}
            yield json.dumps(data, ensure_ascii=False)
            # print("data_res", data_res)
        except Exception as e:
            # pass
            # yield json.dumps({}, ensure_ascii=False)
            print("error", e)
            data = {
                'data': {},
                'event': 'error',
                'dialog_uid': dialog_uid,
            }
            yield json.dumps(data, ensure_ascii=False)
            # print("data_res", data_res)

        ai_msg_data = {
            'dialog_id': dialog_ins_id,
            'role': 'assistant',
            'content': data_res
        }
        create_dialog_history(ai_msg_data, db)

    return EventSourceResponse(general_generate())




